from __future__ import annotations

import os
from pathlib import Path
from typing import Any, Dict, List
import io
import uuid
from datetime import datetime, timezone
import mimetypes
from dotenv import load_dotenv

import requests
from flask import Flask, Response, jsonify, request, send_from_directory, make_response
from werkzeug.utils import secure_filename
try:
    # Optional imports for richer extraction
    from pdfminer.high_level import extract_text as pdf_extract_text  # type: ignore
except Exception:  # pragma: no cover
    pdf_extract_text = None  # type: ignore
try:
    import docx  # type: ignore
except Exception:  # pragma: no cover
    docx = None  # type: ignore
try:
    import pytesseract  # type: ignore
    from PIL import Image  # type: ignore
except Exception:  # pragma: no cover
    pytesseract = None  # type: ignore
    Image = None  # type: ignore

import difflib


BASE_DIR = Path(__file__).resolve().parent

load_dotenv(dotenv_path=BASE_DIR / ".env")

app = Flask(
    __name__,
    static_url_path="",
    static_folder=str(BASE_DIR),
)

# uploads
UPLOAD_DIR = BASE_DIR / "uploads"
UPLOAD_DIR.mkdir(exist_ok=True)
app.config["MAX_CONTENT_LENGTH"] = 16 * 1024 * 1024  # 16MB

# in-memory session logs: { sid: [ {ts, role, content, attachments?, extracted?} ] }
SESSION_LOGS: Dict[str, List[Dict[str, Any]]] = {}

def _now_iso() -> str:
    return datetime.now(timezone.utc).astimezone().isoformat(timespec="seconds")

def _get_sid() -> str:
    sid = request.cookies.get("CHATBOX_SID")
    if sid and isinstance(sid, str) and len(sid) <= 64:
        return sid
    # fallback ephemeral id (not persisted unless set at /)
    return uuid.uuid4().hex


# ---- Teaching prompt helpers ----
TOPIC_NAME_MAP: Dict[str, str] = {
    "c-id": "Identidades (Salud y bienestar; Creencias y valores; Subculturas; Lengua e identidad)",
    "c-os": "Organización Social (Relaciones sociales; Comunidad; Participación social; Educación; El mundo laboral)",
    "c-ih": "Ingenio Humano (Entretenimiento; Expresiones artísticas; Comunicación y medios; Tecnología; Innovación científica)",
    "c-cp": "Compartir el Planeta (Medio ambiente; Derechos humanos; Paz y conflictos; Globalización; Ética)",
    "c-ex": "Experiencias (Actividades recreativas; Vacaciones y viajes; Historias de vida; Costumbres y tradiciones; Migración)",
}

TOPIC_SCOPE_MAP: Dict[str, str] = {
    "c-id": "Identidades",
    "c-os": "Organización Social",
    "c-ih": "Ingenio Humano",
    "c-cp": "Compartir el Planeta",
    "c-ex": "Experiencias",
}

# 子话题种子，供首问多样化
SUBTOPIC_SEEDS: Dict[str, List[str]] = {
    "c-id": [
        "Salud y bienestar",
        "Creencias y valores",
        "Subculturas",
        "Lengua e identidad",
    ],
    "c-os": [
        "Relaciones sociales",
        "Comunidad",
        "Participación social",
        "Educación",
        "El mundo laboral",
    ],
    "c-ih": [
        "Entretenimiento",
        "Expresiones artísticas",
        "Comunicación y medios",
        "Tecnología",
        "Innovación científica",
    ],
    "c-cp": [
        "Medio ambiente",
        "Derechos humanos",
        "Paz y conflictos",
        "Globalización",
        "Ética",
    ],
    "c-ex": [
        "Actividades recreativas",
        "Vacaciones y viajes",
        "Historias de vida",
        "Costumbres y tradiciones",
        "Migración",
    ],
}

# 记录各主题最近的开场问题，避免重复
LAST_OPENING_QUESTIONS: Dict[str, List[str]] = {}

def compose_teacher_system_prompt(topic_key: str | None) -> str:
    topic_scope = TOPIC_SCOPE_MAP.get(str(topic_key or ""), "")
    topic_details = TOPIC_NAME_MAP.get(str(topic_key or ""), "(tema libre)")
    # 合并用户提供的模板，明确输出规范
    return (
        "Eres un profesor universitario de español (nativo) amable y motivador. "
        "Tema principal: " + topic_scope + ". "
        "Ámbito específico: " + topic_details + ". "
        "El estudiante es principiante (IB Ab Initio).\n\n"
        "Instrucciones didácticas:\n"
        "- Haz SIEMPRE UNA sola pregunta abierta por turno, con tiempos verbales dentro de: presente / pretérito indefinido / futuro / imperfecto (ajusta según el nivel).\n"
        "- IMPORTANTE: Todas las preguntas y conversaciones DEBEN girar exclusivamente alrededor del tema principal y sus subtemas especificados.\n"
        "- Si el estudiante YA ha escrito algo en su último turno, DEBES evaluar su respuesta siguiendo el esquema A–F. Solo cuando NO haya ninguna respuesta del estudiante, puedes dejar A–C sin puntuar.\n\n"
        "Formato de salida OBLIGATORIO (usa exactamente estos encabezados en negrita; responde en español; no añadas texto fuera del esquema):\n"
        "1) **Pregunta**: (una sola pregunta abierta relacionada con el tema).\n"
        "2) **A. Análisis de errores**: especifica tipo (gramática/ortografía) y cuenta total (p. ej., '2 de gramática, 1 de ortografía').\n"
        "3) **B. Evaluación de la calidad lingüística**: gramática y vocabulario (diversidad, complejidad, precisión), repetición, relevancia/coherencia, comprensibilidad.\n"
        "4) **C. Evaluación de la competencia** (escala total de 7 puntos; usa una decimal, p. ej., 2.3/7.0; puntúa siempre según el desempeño; NUNCA asignes 0; si la calidad es muy baja usa como mínimo 0.5/7.0):\n   - Gramática: X.X/7.0\n   - Léxico: X.X/7.0\n   - Contenido: X.X/7.0\n   - **Promedio**: X.X/7.0\n"
        "5) **D. Vocabulario nuevo**: 2–3 ítems en la forma 'palabra/expresión → English'.\n"
        "6) **E. Pregunta de seguimiento**: una sola pregunta nueva relacionada con el tema principal.\n"
        "7) **F. Enlace cultural**: un dato cultural breve de un país hispanohablante relacionado con la respuesta del estudiante y el tema principal.\n"
        "Respeta el esquema incluso si el estudiante aún no ha respondido (en ese caso, deja vacías A–C con una nota breve y ofrece la pregunta en 1)."
    )


def _enforce_sections(reply: str) -> str:
    """Si la respuesta no contiene todas las secciones requeridas, pedir al modelo que la reformatee.

    Retorna el texto final (reformateado si fue necesario).
    """
    required = [
        "**Pregunta**",
        "**A. Análisis de errores**",
        "**B. Evaluación de la calidad lingüística**",
        "**C. Evaluación de la competencia**",
        "**D. Vocabulario nuevo**",
        "**E. Pregunta de seguimiento**",
        "**F. Enlace cultural**",
    ]
    if all(s in reply for s in required):
        return reply

    api_key_env = os.environ.get("DEEPSEEK_API_KEY", "")
    if not api_key_env:
        return reply

    try:
        payload = {
            "model": "deepseek-chat",
            "messages": [
                {
                    "role": "system",
                    "content": (
                        "Reformatea el texto del usuario EXACTAMENTE al siguiente esquema en español: "
                        "1) **Pregunta**; 2) **A. Análisis de errores**; 3) **B. Evaluación de la calidad lingüística**; "
                        "4) **C. Evaluación de la competencia** con tres subnotas y promedio (escala total 7 puntos, una decimal, p. ej., 2.3/7.0; NUNCA 0; mínimo 0.5/7.0); 5) **D. Vocabulario nuevo** (2–3 ítems 'español → English'); "
                        "6) **E. Pregunta de seguimiento**; 7) **F. Enlace cultural**. No añadas texto fuera del esquema."
                    ),
                },
                {"role": "user", "content": reply},
            ],
            "temperature": 0.2,
            "max_tokens": 800,
        }
        resp = requests.post(
            "https://api.deepseek.com/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {api_key_env}",
                "Content-Type": "application/json",
            },
            json=payload,
            timeout=40,
        )
        if resp.status_code < 300:
            j = resp.json()
            new_reply = (
                (j.get("choices") or [{}])[0].get("message", {}).get("content", "").strip()
            )
            if new_reply:
                return new_reply
    except Exception:
        pass
    return reply


def _safe_read_text(path: Path, max_chars: int = 8000) -> str:
    """Extract text from supported file types, clipped to max_chars.

    Supports: .txt/.md/.log/.csv/.json/.html as plain text; .pdf via pdfminer; .docx via python-docx; images via OCR.
    Returns empty string if unsupported or on failure.
    """
    try:
        suffix = path.suffix.lower()
        if suffix in {".txt", ".md", ".log", ".csv", ".json", ".html", ".htm"}:
            data = path.read_bytes()
            text = data.decode("utf-8", errors="replace")
            return text[:max_chars]

        if suffix == ".pdf" and pdf_extract_text is not None:
            try:
                text = pdf_extract_text(str(path)) or ""
                return text[:max_chars]
            except Exception:
                return ""

        if suffix == ".docx" and docx is not None:
            try:
                d = docx.Document(str(path))
                text = "\n".join(p.text for p in d.paragraphs)
                return text[:max_chars]
            except Exception:
                return ""

        # Image OCR
        if suffix in {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"} and pytesseract is not None and Image is not None:
            try:
                img = Image.open(str(path))
                # OCR with multiple languages for better accuracy
                text = pytesseract.image_to_string(img, lang='eng+spa', config='--psm 6')
                return text[:max_chars]
            except Exception:
                return ""

        # Fallback: try mime guess and read text/*
        mtype, _ = mimetypes.guess_type(str(path))
        if mtype and mtype.startswith("text/"):
            data = path.read_bytes()
            text = data.decode("utf-8", errors="replace")
            return text[:max_chars]
    except Exception:
        return ""
    return ""


@app.get("/")
def root() -> Response:
    resp = make_response(send_from_directory(str(BASE_DIR), "chatbot.html"))
    if not request.cookies.get("CHATBOX_SID"):
        resp.set_cookie("CHATBOX_SID", uuid.uuid4().hex, max_age=30*24*3600, httponly=True, samesite="Lax")
    return resp


@app.post("/api/ask")
def ask() -> Response:
    data: Dict[str, Any] = request.get_json(silent=True) or {}
    tag = str(data.get("tag", "")).strip()
    if not tag:
        return Response("Missing 'tag'", status=400)

    api_key = os.environ.get("DEEPSEEK_API_KEY")
    if not api_key:
        # 让前端可见明确错误
        return Response("DEEPSEEK_API_KEY not set", status=500)

    try:
        def is_diverse(qs: List[str], cand: str, threshold: float = 0.82) -> bool:
            cand_n = cand.strip().lower()
            for q in qs:
                if difflib.SequenceMatcher(None, q.strip().lower(), cand_n).ratio() >= threshold:
                    return False
            return True

        used = LAST_OPENING_QUESTIONS.get(tag, [])

        def build_payload(seed: str | None = None):
            user_content = (
                f"Tema principal: {TOPIC_SCOPE_MAP.get(tag, '')}. "
                f"Ámbito específico: {TOPIC_NAME_MAP.get(tag, '')}. "
            )
            if seed:
                user_content += f"Subtema sugerido: {seed}. "
            user_content += "Redacta UNA sola pregunta abierta, natural y diferente a las anteriores, en español."
            return {
                "model": "deepseek-chat",
                "messages": [
                    {
                        "role": "system",
                        "content": (
                            "Eres un profesor de español. Genera SOLO UNA pregunta abierta, breve y clara, en español, "
                            "para iniciar una conversación dentro del tema indicado. No des explicaciones ni extras; "
                            "devuelve solo la pregunta."
                        ),
                    },
                    {"role": "user", "content": user_content},
                ],
                "temperature": 0.95,
                "max_tokens": 128,
            }

        seeds = SUBTOPIC_SEEDS.get(tag, [None])
        attempts = 0
        question = ""
        while attempts < min(3, max(1, len(seeds))):
            seed = seeds[attempts % len(seeds)] if seeds else None
            payload = build_payload(seed if isinstance(seed, str) else None)
            resp = requests.post(
                "https://api.deepseek.com/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json=payload,
                timeout=30,
            )
            if resp.status_code >= 300:
                attempts += 1
                continue
            j = resp.json()
            cand = (
                (j.get("choices") or [{}])[0]
                .get("message", {})
                .get("content", "")
                .strip()
            )
            if cand and is_diverse(used, cand):
                question = cand
                break
            attempts += 1

        if not question:
            # fallback 再取一次不带种子
            payload = build_payload(None)
            resp = requests.post(
                "https://api.deepseek.com/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {api_key}",
                    "Content-Type": "application/json",
                },
                json=payload,
                timeout=30,
            )
            if resp.status_code < 300:
                j = resp.json()
                question = (
                    (j.get("choices") or [{}])[0]
                    .get("message", {})
                    .get("content", "")
                    .strip()
                )

        if question:
            # 写入缓存，保留最近 10 条
            bucket = LAST_OPENING_QUESTIONS.setdefault(tag, [])
            bucket.append(question)
            if len(bucket) > 10:
                del bucket[:-10]
            return jsonify({"question": question})
        return Response("No question generated", status=502)
    except Exception as exc:  # noqa: BLE001
        return Response(f"Error: {exc}", status=500)


@app.post("/api/chat")
def chat() -> Response:
    data: Dict[str, Any] = request.get_json(silent=True) or {}
    messages = data.get("messages")
    attachments: List[Dict[str, Any]] = data.get("attachments") or []
    topic_key = str(data.get("topicKey", ""))

    if not isinstance(messages, list) or not messages:
        return Response("Missing 'messages' (list)", status=400)

    try:
        # find last user message for scoring fix
        last_user = ""
        if isinstance(messages, list):
            for m in reversed(messages):
                if isinstance(m, dict) and m.get("role") == "user":
                    last_user = str(m.get("content", ""))
                    break
        
        # 读取 API Key（用于下游请求）
        api_key = os.environ.get("DEEPSEEK_API_KEY")
        if not api_key:
            return Response("DEEPSEEK_API_KEY not set", status=500)
        
        # Build attachments context
        context_chunks: List[str] = []
        for att in attachments:
            # prefer url, else name
            url: str = str(att.get("url", ""))
            name: str = str(att.get("name", ""))
            # Map url like /uploads/<filename> to local path
            filename = name
            if url.startswith("/uploads/"):
                filename = url.split("/uploads/", 1)[-1]
            local_path = (UPLOAD_DIR / filename).resolve()
            # Security: ensure path is inside UPLOAD_DIR
            if UPLOAD_DIR not in local_path.parents and local_path != UPLOAD_DIR:
                continue
            if not local_path.exists() or not local_path.is_file():
                continue
            text = _safe_read_text(local_path)
            if text:
                context_chunks.append(f"Archivo: {filename}\n\n{text}")

        context_text = "\n\n---\n\n".join(context_chunks)

        llm_messages: List[Dict[str, str]] = []
        # Teacher system prompt with topic
        llm_messages.append({
            "role": "system",
            "content": compose_teacher_system_prompt(topic_key),
        })
        if context_text:
            llm_messages.append({
                "role": "system",
                "content": (
                    "Usa el siguiente contexto de archivos adjuntos cuando sea relevante. "
                    "Si la pregunta no está cubierta por el contexto, responde normalmente.\n\n"
                    f"Contexto:\n{context_text}"
                ),
            })
        # then original messages
        if isinstance(messages, list):
            llm_messages.extend(messages)
        else:
            llm_messages.extend([])
        payload = {
            "model": "deepseek-chat",
            "messages": llm_messages,
            "temperature": float(data.get("temperature", 0.7)),
            "max_tokens": int(data.get("max_tokens", 512)),
        }
        resp = requests.post(
            "https://api.deepseek.com/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json=payload,
            timeout=60,
        )
        if resp.status_code >= 300:
            return Response(f"DeepSeek error: {resp.status_code} {resp.text}", status=502)
        j = resp.json()
        reply = (
            (j.get("choices") or [{}])[0]
            .get("message", {})
            .get("content", "")
            .strip()
        )
        # Enforce sections A–F if missing
        reply = _enforce_sections(reply)

        # 如果用户已有回答，但模型仍返回占位或0分，进行最低分修正
        if last_user.strip():
            try:
                if "Gramática:" in reply and ("-/7.0" in reply or "/7.0\nD." in reply or " 0/7.0" in reply):
                    # 简单正则式替换到 0.5/7.0
                    reply = reply.replace("Gramática: -/7.0", "Gramática: 0.5/7.0")
                    reply = reply.replace("Léxico: -/7.0", "Léxico: 0.5/7.0")
                    reply = reply.replace("Contenido: -/7.0", "Contenido: 0.5/7.0")
                    # 若平均分仍为占位，按0.5计算
                    reply = reply.replace("Promedio: -/7.0", "Promedio: 0.5/7.0")
                    reply = reply.replace(" 0/7.0", " 0.5/7.0")
            except Exception:
                pass
        # session logging
        sid = _get_sid()
        logs = SESSION_LOGS.setdefault(sid, [])
        # best-effort: last user text from messages
        logs.append({
            "ts": _now_iso(), "role": "user", "content": last_user, "attachments": attachments
        })
        logs.append({
            "ts": _now_iso(), "role": "assistant", "content": reply, "extracted": context_text
        })
        return jsonify({"reply": reply})
    except Exception as exc:  # noqa: BLE001
        return Response(f"Error: {exc}", status=500)


@app.post("/api/translate")
def translate() -> Response:
    data: Dict[str, Any] = request.get_json(silent=True) or {}
    text = str(data.get("text", "")).strip()
    target = str(data.get("target", "es")).strip().lower()
    if not text:
        return Response("Missing 'text'", status=400)

    api_key = os.environ.get("DEEPSEEK_API_KEY")
    if not api_key:
        return Response("DEEPSEEK_API_KEY not set", status=500)

    # decide prompt
    if target == "en":
        sys_prompt = "You are a translator. Translate the given text into English. Return only the translated text with no explanations."
    elif target == "es":
        sys_prompt = "Eres un traductor. Traduce cualquier texto dado al español. Devuelve solo el texto traducido, sin explicaciones."
    else:
        sys_prompt = f"You are a translator. Translate the given text into {target}. Return only the translated text with no explanations."

    try:
        payload = {
            "model": "deepseek-chat",
            "messages": [
                {"role": "system", "content": sys_prompt},
                {"role": "user", "content": text},
            ],
            "temperature": 0.2,
            "max_tokens": 256,
        }
        resp = requests.post(
            "https://api.deepseek.com/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json=payload,
            timeout=30,
        )
        if resp.status_code >= 300:
            return Response(f"DeepSeek error: {resp.status_code} {resp.text}", status=502)
        j = resp.json()
        translated = (
            (j.get("choices") or [{}])[0]
            .get("message", {})
            .get("content", "")
            .strip()
        )
        # backward compatible key for Spanish
        if target == "es":
            return jsonify({"translated": translated, "spanish": translated, "target": target})
        return jsonify({"translated": translated, "target": target})
    except Exception as exc:  # noqa: BLE001
        return Response(f"Error: {exc}", status=500)


@app.post("/api/upload")
def upload() -> Response:
    if "file" not in request.files:
        return Response("Missing file", status=400)
    f = request.files["file"]
    if not f or f.filename == "":
        return Response("Empty filename", status=400)

    filename = secure_filename(f.filename)
    save_path = UPLOAD_DIR / filename
    # prevent overwrite by appending index
    idx = 1
    stem = save_path.stem
    suffix = save_path.suffix
    while save_path.exists():
        filename = f"{stem}_{idx}{suffix}"
        save_path = UPLOAD_DIR / filename
        idx += 1
    f.save(str(save_path))
    url = f"/uploads/{filename}"
    return jsonify({"filename": filename, "url": url})


@app.get("/uploads/<path:filename>")
def serve_upload(filename: str) -> Response:
    return send_from_directory(str(UPLOAD_DIR), filename)


@app.get("/api/export_pdf")
def export_pdf() -> Response:
    sid = _get_sid()
    items = SESSION_LOGS.get(sid, [])
    if not items:
        return Response("No session logs", status=404)

    try:
        # Lazy import fpdf to speed normal requests
        from fpdf import FPDF  # type: ignore

        pdf = FPDF(format="A4")
        pdf.set_auto_page_break(auto=True, margin=12)
        pdf.add_page()
        pdf.set_font("Helvetica", size=12)

        pdf.cell(0, 8, txt="Chat export", ln=1)
        pdf.set_font_size(10)
        pdf.set_text_color(100, 100, 100)
        export_time = _now_iso()
        pdf.cell(0, 6, txt=f"Exported at: {export_time}", ln=1)
        pdf.cell(0, 6, txt=f"Total messages: {len(items)}", ln=1)
        pdf.set_text_color(0, 0, 0)

        for it in items:
            ts = str(it.get("ts", ""))
            role = str(it.get("role", ""))
            content = str(it.get("content", ""))
            extracted = str(it.get("extracted", ""))
            attachments = it.get("attachments") or []

            # header line
            pdf.set_font("Helvetica", style="B", size=11)
            hdr = f"[{ts}] {role.capitalize()}"
            pdf.multi_cell(0, 6, txt=hdr)
            # content
            pdf.set_font("Helvetica", size=11)
            pdf.multi_cell(0, 6, txt=content)
            # attachments names
            if attachments:
                names = ", ".join(str(a.get("name")) for a in attachments if isinstance(a, dict) and a.get("name"))
                if names:
                    pdf.set_text_color(100, 100, 100)
                    pdf.set_font_size(10)
                    pdf.multi_cell(0, 5, txt=f"Attachments: {names}")
                    pdf.set_text_color(0, 0, 0)
                    pdf.set_font_size(11)
            # extracted context
            if extracted:
                pdf.set_text_color(120, 120, 120)
                pdf.set_font_size(10)
                pdf.multi_cell(0, 5, txt=f"Extracted:\n{extracted}")
                pdf.set_text_color(0, 0, 0)
                pdf.set_font_size(11)
            pdf.ln(2)

        buf = io.BytesIO()
        pdf.output(buf)
        data = buf.getvalue()
        return Response(data, mimetype="application/pdf", headers={
            "Content-Disposition": f"attachment; filename=chat-export-{datetime.now().strftime('%Y%m%d-%H%M%S')}.pdf"
        })
    except Exception as exc:  # noqa: BLE001
        return Response(f"Error generating PDF: {exc}", status=500)


if __name__ == "__main__":
    port = int(os.environ.get("PORT", "8008"))
    app.run(host="127.0.0.1", port=port, debug=True)


